from typing import Dict, Any
from agents.components.base import BaseComponent, ComponentConfig, ComponentType
from agents.components.llm.ollama_llm import OllamaLLM

class OllamaComponent(BaseComponent):
    """Ollama模型组件实现"""
    
    def __init__(self, config: ComponentConfig):
        super().__init__(config)
        self.model = OllamaLLM(self.config.config)
    
    @classmethod
    def get_component_type(cls) -> ComponentType:
        return ComponentType.LLM
    
    @classmethod
    def get_component_name(cls) -> str:
        return "Ollama模型"
    
    @classmethod
    def get_component_description(cls) -> str:
        return "使用Ollama的API进行模型调用"
    
    @classmethod
    def get_config_schema(cls) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "system_prompt": {
                    "type": "string",
                    "title": "系统提示词",
                    "description": "模型的系统提示词"
                },
                "user_prompt": {
                    "type": "string",
                    "title": "用户提示词",
                    "description": "模型的用户提示词，可以使用{变量}的形式"
                },
                "model_name": {
                    "type": "string",
                    "title": "模型名称",
                    "description": "使用的Ollama模型名称",
                    "default": "llama2"
                },
                "temperature": {
                    "type": "number",
                    "title": "温度",
                    "description": "控制输出的随机性，值越大越随机",
                    "default": 0.7
                },
                "top_p": {
                    "type": "number",
                    "title": "Top P",
                    "description": "控制输出的多样性",
                    "default": 0.9
                },
                "top_k": {
                    "type": "integer",
                    "title": "Top K",
                    "description": "控制每次采样的候选词数量",
                    "default": 40
                },
                "repeat_penalty": {
                    "type": "number",
                    "title": "重复惩罚",
                    "description": "控制重复内容的惩罚力度",
                    "default": 1.1
                },
                "num_ctx": {
                    "type": "integer",
                    "title": "上下文长度",
                    "description": "控制模型处理的上下文长度",
                    "default": 2048
                },
                "stream": {
                    "type": "boolean",
                    "title": "流式输出",
                    "description": "是否使用流式输出",
                    "default": False
                }
            },
            "required": ["model_name"]
        }
    
    def get_input_schema(self) -> Dict[str, Any]:
        return self.config.input_schema
    
    def get_output_schema(self) -> Dict[str, Any]:
        return self.config.output_schema
    
    async def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """执行Ollama模型调用"""
        system_prompt = self.config.config.get("system_prompt", "")
        user_prompt = self.config.config.get("user_prompt", "").format(**input_data)
        
        # 从配置中获取模型参数
        model_params = {
            "temperature": self.config.config.get("temperature"),
            "top_p": self.config.config.get("top_p"),
            "top_k": self.config.config.get("top_k"),
            "repeat_penalty": self.config.config.get("repeat_penalty"),
            "num_ctx": self.config.config.get("num_ctx")
        }
        
        # 过滤掉None值
        model_params = {k: v for k, v in model_params.items() if v is not None}
        
        response = self.model.chat(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            **model_params
        )
        
        return {"response": response}
    
    def validate_input(self, input_data: Dict[str, Any]) -> bool:
        """验证输入数据"""
        required_fields = self.config.input_schema.get("required", [])
        return all(field in input_data for field in required_fields)
    
    def validate_output(self, output_data: Dict[str, Any]) -> bool:
        """验证输出数据"""
        return "response" in output_data